import cv2
import torch
import requests
import numpy as np

# **ESP32-CAM 视频流 URL**
ESP32_CAM_URL = "http://192.168.43.15:81/stream"  # 请确保 IP 地址正确

# **YOLOv5 模型路径**
YOLO_MODEL_PATH = "E:\\yolov5-6.0\\best.pt"  # 你的 YOLOv5 训练模型

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# **加载 YOLOv5 模型**
model = torch.hub.load('.', 'custom', path=YOLO_MODEL_PATH, source='local', device=device)
model.eval()

# **打开 ESP32-CAM 视频流**
cap = cv2.VideoCapture(ESP32_CAM_URL)
if not cap.isOpened():
    print("❌ 无法连接 ESP32-CAM 视频流")
    exit()

frame_count = 0
while True:
    ret, frame = cap.read()  # 读取视频帧
    if not ret:
        print("❌ 读取失败，尝试重新连接")
        cap = cv2.VideoCapture(ESP32_CAM_URL)
        continue

    frame_count += 1
    if frame_count % 2 == 0:  # 每两帧处理一次
        results = model(frame)
        # 处理检验

    # 运行 YOLOv5 目标检测
    results = model(frame)



    # 绘制检测结果
    for result in results.xyxy[0]:  # 遍历检测到的目标
        x1, y1, x2, y2, conf, cls = result
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        label = f"{int(cls)}: {conf:.2f}"
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 画框
        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 显示检测结果
    cv2.imshow("ESP32-CAM YOLOv5 Detection", frame)

    # 按 `q` 退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()
